#!/bin/bash

# 检测系统发行版
if grep -q "Ubuntu 20.04" /etc/os-release; then
    echo "Detected Ubuntu 20.04"
    #ubuntu 20.04
    
    set -e
    
    # 创建 /opt 目录
    echo "cd /opt/"
    mkdir -p /opt/
    cd /opt/
    
    # 添加 MongoDB 4.4 的软件源
    echo "deb [ arch=amd64,arm64 signed-by=/usr/share/keyrings/mongodb-server-4.4.gpg ] https://repo.mongodb.org/apt/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME")/mongodb-org/4.4 multiverse" | tee /etc/apt/sources.list.d/mongodb-org-4.4.list
    
    # 检查并安装 curl
    if ! command -v curl &> /dev/null; then
        echo "install curl ..."
        apt-get install curl -y
    fi
    
    # 导入 MongoDB GPG 密钥
    curl -fsSL https://www.mongodb.org/static/pgp/server-4.4.asc | gpg -o /usr/share/keyrings/mongodb-server-4.4.gpg --dearmor
    
    # 更新软件包列表
    echo "Updating package list..."
    apt-get update -y
    apt install software-properties-common
    add-apt-repository ppa:deadsnakes/ppa
    apt update -y && apt-get upgrade -y
    
    # 安装依赖
    echo "Installing dependencies..."
    apt-get install -y python3.6 mongodb-org rabbitmq-server python3.6-dev g++ git nginx fontconfig \
    unzip wget gnupg lsb-release python3.6-distutils -y
    
    # 安装 pip3.6
    if [ ! -f /usr/local/bin/pip3.6 ]; then
        echo "install pip3.6"
        curl https://raw.gitcode.com/msmoshang/arl_files/blobs/fd48c7fdef802d8bb86ace74134c553f0317258c/get-pip.py -o get-pip.py
        # 国内下载失败可以打开
        python3.6 get-pip.py -i https://mirrors.aliyun.com/pypi/simple/
        pip3.6 --version
        pip3.6 config set global.index-url https://mirrors.aliyun.com/pypi/simple/
        rm -rf get-pip.py
    fi
    
    # 检查并安装 alien
    if ! command -v alien &> /dev/null; then
        echo "install alien..."
        apt-get install alien -y
    fi
    
    # 安装 nmap
    if ! command -v nmap &> /dev/null; then
        echo "Installing nmap..."
        apt install -y nmap
    fi
    
    
    if [ ! -d arl-extension ]; then
        echo "git clone arl-extension"
        git clone https://gitee.com/lonelywolf09/arl-extension.git
    fi
    
    
    if ! command -v nuclei &> /dev/null
    then
        echo "install nuclei_3.3.5 ..."
        cd arl-extension/
        unzip nuclei_3.3.5_linux_amd64.zip && cp nuclei /usr/bin/ && rm -f nuclei_3.3.5_linux_amd64.zip
        nuclei -ut
    fi
    
    
    if ! command -v wih &> /dev/null
    then
        echo "install wih ..."
        ## 安装 WIH
        cd wih/ && cp wih_linux_amd64 /usr/bin/wih && chmod +x /usr/bin/wih
    fi
    
    # 启动服务
    echo "Starting services..."
    systemctl enable mongod
    systemctl start mongod
    systemctl enable rabbitmq-server
    systemctl start rabbitmq-server

    cd /opt
    
    # 克隆项目
    if [ ! -d ARL ]; then
        echo "Cloning ARL project..."
        git clone https://gitee.com/lonelywolf09/ARL.git
    fi
    
    if [ ! -d "ARL-NPOC" ]; then
        echo "Cloning ARL-NPOC project..."
        git clone https://gitee.com/lonelywolf09/ARL-NPOC.git
    fi
    
    cd ARL-NPOC
    echo "Installing POC requirements..."
    pip3.6 install --user -r requirements.txt
    pip3.6 install -e .
    

    cd /opt/arl-extension
    
    if [ ! -f /usr/local/bin/ncrack ]; then
        echo "install ncrack ..."
        cp ncrack /usr/local/bin/ncrack &&  chmod +x /usr/local/bin/ncrack

    fi
    
    mkdir -p /usr/local/share/ncrack
    if [ ! -f /usr/local/share/ncrack/ncrack-services ]; then
        echo "install ncrack-services ..."
        cp ncrack-services /usr/local/share/ncrack/ncrack-services
    fi
    
    mkdir -p /data/GeoLite2
    if [ ! -f /data/GeoLite2/GeoLite2-ASN.mmdb ]; then
        echo "install GeoLite2-ASN.mmdb ..."
        cp GeoLite2-ASN/GeoLite2-ASN.mmdb /data/GeoLite2/GeoLite2-ASN.mmdb
    fi
    
    if [ ! -f /data/GeoLite2/GeoLite2-City.mmdb ]; then
        echo "install GeoLite2-City.mmdb ..."
        cp GeoLite2-City/GeoLite2-City.mmdb /data/GeoLite2/GeoLite2-City.mmdb

    fi
    
    cd /opt/ARL
    
    # 配置 RabbitMQ 用户
    if [ ! -f rabbitmq_user ]; then
        echo "Adding RabbitMQ user..."
        rabbitmqctl add_user arl arlpassword
        rabbitmqctl add_vhost arlv2host
        rabbitmqctl set_user_tags arl arltag
        rabbitmqctl set_permissions -p arlv2host arl ".*" ".*" ".*"
        echo "Initializing ARL user..."
        mongo 127.0.0.1:27017/arl docker/mongo-init.js
        touch rabbitmq_user
    fi
    
    # 安装 ARL 的 Python 依赖
    echo "Installing ARL requirements..."
    pip3.6 install -r requirements.txt
    
    # 创建配置文件
    if [ ! -f app/config.yaml ]; then
        echo "Creating config.yaml"
        cp app/config.yaml.example app/config.yaml
    fi
    
    # 设置 phantomjs
    if [ ! -f /usr/local/bin/phantomjs ]; then
        echo "Installing phantomjs..."
        ln -s `pwd`/app/tools/phantomjs /usr/local/bin/phantomjs
    fi
    
    
    #check nginx log format
    # 定义路径
    NGINX_CONF="/etc/nginx/nginx.conf"
    NEW_CONF="/opt/ARL/misc/nginx.conf"
    
    # 检查 log_format 是否已存在
    if grep -q 'log_format main' "$NGINX_CONF"; then
        echo "log_format already exists in $NGINX_CONF. No changes made."
    else
        echo "log_format not found. Replacing $NGINX_CONF with new configuration..."
        
        # 备份当前配置文件
        cp "$NGINX_CONF" "${NGINX_CONF}.bak"
        
        # 替换配置文件
        cp "$NEW_CONF" "$NGINX_CONF"
    fi
    
    # 配置 nginx
    if [ ! -f /etc/nginx/sites-available/arl.conf ]; then
        echo "Copying arl.conf to nginx configuration..."
        cp misc/arl.conf /etc/nginx/sites-available/arl.conf
        ln -s /etc/nginx/sites-available/arl.conf /etc/nginx/sites-enabled/
    fi
    
    
    
    # 配置 SSL 证书
    if [ ! -f /etc/ssl/certs/dhparam.pem ]; then
        echo "Downloading dhparam.pem..."
        curl https://ssl-config.mozilla.org/ffdhe2048.txt > /etc/ssl/certs/dhparam.pem
    fi
    
    echo "Generating certificate..."
    chmod +x /opt/ARL/docker/worker/gen_crt.sh
    ./docker/worker/gen_crt.sh
    
    # 配置 systemd 服务
    cd /opt/ARL/
    
    for service_file in arl-web.service arl-worker.service arl-worker-github.service arl-scheduler.service; do
        if [ ! -f /etc/systemd/system/$service_file ]; then
            echo "Copying $service_file..."
            cp misc/$service_file /etc/systemd/system/
        fi
    done
    nginx -s reload
    chmod -R +x /opt/ARL/app/tools
    
    # 启动服务
    echo "Starting ARL services..."
    systemctl enable arl-web
    systemctl start arl-web
    systemctl enable arl-worker
    systemctl start arl-worker
    systemctl enable arl-worker-github
    systemctl start arl-worker-github
    systemctl enable arl-scheduler
    systemctl start arl-scheduler
    systemctl enable nginx
    systemctl start nginx
    
    # 检查服务状态
    systemctl status arl-web
    systemctl status arl-worker
    systemctl status arl-worker-github
    systemctl status arl-scheduler

    
    echo "Installation complete!"

    elif grep -q "CentOS Linux 7" /etc/os-release; then
    echo "Detected CentOS 7"
    #centos 7
    set -e
    
    echo "cd /opt/"
    
    mkdir -p /opt/
    cd /opt/
    
tee /etc/yum.repos.d/mongodb-org-4.0.repo <<"EOF"
[mongodb-org-4.0]
name=MongoDB Repository
baseurl=https://repo.mongodb.org/yum/redhat/$releasever/mongodb-org/4.0/x86_64/
gpgcheck=1
enabled=1
gpgkey=https://www.mongodb.org/static/pgp/server-4.0.asc
EOF
    
    
    echo "install dependencies ..."
    yum install epel-release -y
    yum install python36 mongodb-org-server mongodb-org-shell rabbitmq-server python36-devel gcc-c++ git \
    nginx  fontconfig wqy-microhei-fonts unzip wget -y
    
    if [ ! -f /usr/bin/python3.6 ]; then
        echo "link python3.6"
        ln -s /usr/bin/python36 /usr/bin/python3.6
    fi
    
    if [ ! -f /usr/local/bin/pip3.6 ]; then
        echo "install  pip3.6"
        python3.6 -m ensurepip --default-pip
        python3.6 -m pip install --upgrade pip
        pip3.6 --version
    fi
    
    if ! command -v nmap &> /dev/null
    then
        echo "install nmap-7.93-1 ..."
        rpm -vhU https://nmap.org/dist/nmap-7.93-1.x86_64.rpm
    fi


    if [ ! -d arl-extension ]; then
        echo "git clone arl-extension"
        git clone https://gitee.com/lonelywolf09/arl-extension.git
    fi
    
    
    if ! command -v nuclei &> /dev/null
    then
        echo "install nuclei_3.3.5 ..."
        cd arl-extension/
        unzip nuclei_3.3.5_linux_amd64.zip && cp nuclei /usr/bin/ && rm -f nuclei_3.3.5_linux_amd64.zip
        nuclei -ut
    fi
    
    
    if ! command -v wih &> /dev/null
    then
        echo "install wih ..."
        ## 安装 WIH
        cd wih/ && cp wih_linux_amd64 /usr/bin/wih && chmod +x /usr/bin/wih
    fi
    
    
    echo "start services ..."
    systemctl enable mongod
    systemctl start mongod
    systemctl enable rabbitmq-server
    systemctl start rabbitmq-server
    
    cd /opt
    if [ ! -d ARL ]; then
        echo "git clone ARL"
        git clone https://gitee.com/lonelywolf09/ARL.git
    fi
    
    if [ ! -d "ARL-NPOC" ]; then
        echo "git clone ARL-NPOC "
        git clone https://gitee.com/lonelywolf09/ARL-NPOC.git
    fi
    
    cd ARL-NPOC
    echo "install poc requirements ..."
    pip3.6 install --user -r requirements.txt
    pip3.6 install -e .


    cd /opt/arl-extension
    
    if [ ! -f /usr/local/bin/ncrack ]; then
        echo "install ncrack ..."
        cp ncrack /usr/local/bin/ncrack &&  chmod +x /usr/local/bin/ncrack

    fi
    
    mkdir -p /usr/local/share/ncrack
    if [ ! -f /usr/local/share/ncrack/ncrack-services ]; then
        echo "install ncrack-services ..."
        cp ncrack-services /usr/local/share/ncrack/ncrack-services
    fi
    
    mkdir -p /data/GeoLite2
    if [ ! -f /data/GeoLite2/GeoLite2-ASN.mmdb ]; then
        echo "install GeoLite2-ASN.mmdb ..."
        cp GeoLite2-ASN/GeoLite2-ASN.mmdb /data/GeoLite2/GeoLite2-ASN.mmdb
    fi
    
    if [ ! -f /data/GeoLite2/GeoLite2-City.mmdb ]; then
        echo "install GeoLite2-City.mmdb ..."
        cp GeoLite2-City/GeoLite2-City.mmdb /data/GeoLite2/GeoLite2-City.mmdb

    fi
    
    cd /opt/ARL
    
    if [ ! -f rabbitmq_user ]; then
        echo "add rabbitmq user"
        rabbitmqctl add_user arl arlpassword
        rabbitmqctl add_vhost arlv2host
        rabbitmqctl set_user_tags arl arltag
        rabbitmqctl set_permissions -p arlv2host arl ".*" ".*" ".*"
        echo "init arl user"
        mongo 127.0.0.1:27017/arl docker/mongo-init.js
        touch rabbitmq_user
    fi
    
    echo "install arl requirements ..."
    pip3.6 install -r requirements.txt
    if [ ! -f app/config.yaml ]; then
        echo "create config.yaml"
        cp app/config.yaml.example  app/config.yaml
    fi
    
    if [ ! -f /usr/bin/phantomjs ]; then
        echo "install phantomjs"
        ln -s `pwd`/app/tools/phantomjs  /usr/bin/phantomjs
    fi
    
    if [ ! -f /etc/nginx/conf.d/arl.conf ]; then
        echo "copy arl.conf"
        cp misc/arl.conf /etc/nginx/conf.d
    fi
    
    
    
    if [ ! -f /etc/ssl/certs/dhparam.pem ]; then
        echo "download dhparam.pem"
        curl https://ssl-config.mozilla.org/ffdhe2048.txt > /etc/ssl/certs/dhparam.pem
    fi
    
    
    echo "gen cert ..."
    chmod +x /opt/ARL/docker/worker/gen_crt.sh
    ./docker/worker/gen_crt.sh
    
    
    cd /opt/ARL/
    
    
    if [ ! -f /etc/systemd/system/arl-web.service ]; then
        echo  "copy arl-web.service"
        cp misc/arl-web.service /etc/systemd/system/
    fi
    
    if [ ! -f /etc/systemd/system/arl-worker.service ]; then
        echo  "copy arl-worker.service"
        cp misc/arl-worker.service /etc/systemd/system/
    fi
    
    
    if [ ! -f /etc/systemd/system/arl-worker-github.service ]; then
        echo  "copy arl-worker-github.service"
        cp misc/arl-worker-github.service /etc/systemd/system/
    fi
    
    if [ ! -f /etc/systemd/system/arl-scheduler.service ]; then
        echo  "copy arl-scheduler.service"
        cp misc/arl-scheduler.service /etc/systemd/system/
    fi
    
    nginx -s reload
    chmod -R +x /opt/ARL/app/tools
    
    echo "start arl services ..."
    systemctl enable arl-web
    systemctl start arl-web
    systemctl enable arl-worker
    systemctl start arl-worker
    systemctl enable arl-worker-github
    systemctl start arl-worker-github
    systemctl enable arl-scheduler
    systemctl start arl-scheduler
    systemctl enable nginx
    systemctl start nginx
    
    systemctl status arl-web
    systemctl status arl-worker
    systemctl status arl-worker-github
    systemctl status arl-scheduler

    
    echo "install done"
else
    echo "Unsupported OS version"
    exit 1
fi

# 公共命令，适用于两个系统
# 示例：执行其他通用命令
echo "Executing common commands for both systems"
